#!/usr/bin/env python3
"""统一修复ChromaDB路径配置问题"""

import os
import sys
import yaml
import shutil
from pathlib import Path

def main():
    """统一修复ChromaDB路径配置"""
    print("=== ChromaDB路径统一修复工具 ===")
    
    # 1. 确认项目根目录
    project_root = Path("d:/MCP/mcp2")
    print(f"项目根目录: {project_root}")
    
    # 2. 检查当前工作目录
    current_dir = Path.cwd()
    print(f"当前工作目录: {current_dir}")
    
    # 3. 检查配置文件
    config_file = project_root / "config" / "database.yaml"
    print(f"\n检查配置文件: {config_file}")
    
    if config_file.exists():
        with open(config_file, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        chromadb_config = config.get('databases', {}).get('chromadb', {})
        current_persist_dir = chromadb_config.get('persist_directory', './data/chroma_db')
        print(f"当前配置的persist_directory: {current_persist_dir}")
        
        # 4. 检查可能的ChromaDB数据位置
        possible_paths = [
            project_root / "data" / "chroma_db",
            project_root / "src" / "chroma_db",
            Path("./data/chroma_db").resolve(),
            Path("./src/chroma_db").resolve()
        ]
        
        print("\n检查可能的ChromaDB数据位置:")
        data_locations = []
        for path in possible_paths:
            if path.exists():
                items = list(path.iterdir())
                print(f"✓ {path} (包含 {len(items)} 个项目)")
                if len(items) > 0:
                    data_locations.append((path, len(items)))
            else:
                print(f"✗ {path} (不存在)")
        
        # 5. 确定标准路径
        standard_path = project_root / "data" / "chroma_db"
        print(f"\n标准路径: {standard_path}")
        
        # 6. 如果标准路径不存在或为空，寻找有数据的路径
        if not standard_path.exists() or len(list(standard_path.iterdir())) == 0:
            if data_locations:
                # 找到有数据的路径，移动到标准位置
                source_path, item_count = max(data_locations, key=lambda x: x[1])
                print(f"\n发现数据位置: {source_path} (包含 {item_count} 个项目)")
                
                if source_path != standard_path:
                    print(f"移动数据从 {source_path} 到 {standard_path}")
                    
                    # 确保目标目录存在
                    standard_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # 移动数据
                    if standard_path.exists():
                        shutil.rmtree(standard_path)
                    shutil.move(str(source_path), str(standard_path))
                    print("✓ 数据移动完成")
            else:
                print("\n未找到现有ChromaDB数据，将创建新的标准路径")
                standard_path.mkdir(parents=True, exist_ok=True)
        
        # 7. 更新配置文件
        new_persist_dir = "./data/chroma_db"
        if current_persist_dir != new_persist_dir:
            print(f"\n更新配置文件中的persist_directory: {current_persist_dir} -> {new_persist_dir}")
            chromadb_config['persist_directory'] = new_persist_dir
            
            # 写回配置文件
            with open(config_file, 'w', encoding='utf-8') as f:
                yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
            print("✓ 配置文件已更新")
        else:
            print("\n配置文件中的persist_directory已经是标准路径")
        
        # 8. 验证最终状态
        print("\n=== 最终状态验证 ===")
        print(f"标准路径: {standard_path}")
        print(f"路径存在: {standard_path.exists()}")
        if standard_path.exists():
            items = list(standard_path.iterdir())
            print(f"包含项目: {len(items)}")
            if items:
                print(f"项目列表: {[item.name for item in items[:5]]}{'...' if len(items) > 5 else ''}")
        
        print(f"配置路径: {new_persist_dir}")
        
        # 9. 生成MCP配置建议
        print("\n=== MCP配置建议 ===")
        print("确保MCP配置中的工作目录设置正确:")
        print('{')
        print('  "mcpServers": {')
        print('    "mcp2-stdio": {')
        print('      "command": "python",')
        print('      "args": ["-u", "-m", "mcp.main", "--stdio", "--config", "config"],')
        print('      "cwd": "d:\\\\MCP\\\\mcp2",')
        print('      "env": {')
        print('        "PYTHONPATH": "d:\\\\MCP\\\\mcp2\\\\src"')
        print('      }')
        print('    }')
        print('  }')
        print('}')
        
    else:
        print(f"✗ 配置文件不存在: {config_file}")
    
    print("\n=== 修复完成 ===")

if __name__ == "__main__":
    main()